#include "smartconfig.h"
#include "sc_wifi_mgr.h"
#include "WifiConfig_hal_sys_api.h"
#include <FreeRTOS.h>
#include "task.h"
#include "timers.h"
#include "queue.h"
#include <lwip/netdb.h>
#include <wifi_mgmr_ext.h>
#include <aos/yloop.h>
#include <utils_log.h>
#include <string.h>
#include "WifiConfig_hal_sys_api.h"

#define LOCK_BUFFER_SIZE 12

#define SC_ADDR_FORMATSTRING "%02x:%02x:%02x:%02x:%02x:%02x"
#define SC_MAC_FORMAT(mac) mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]

#define SC_LOCK_TODS   0
#define SC_LOCK_FROMDS 1

enum {
    UDP_CMD_EXIT,
    UDP_CMD_CONNECTED,
    UDP_CMD_SEND,
};

typedef union
{
    uint32_t data;
    struct 
    {
        uint16_t mac_tail;
        uint16_t len;
    };
} __attribute__((packed)) buffer_t;


typedef struct {
    char mac[6];
    char ssid[32];
    char bssid[6];
    int32_t state;
    int32_t channel;
    uint16_t offset_tods;
    uint16_t offset_frds;
    uint16_t seqnum_tods;
    uint16_t seqnum_frds;
    buffer_t buffer_tods[LOCK_BUFFER_SIZE];
    buffer_t buffer_frds[LOCK_BUFFER_SIZE];
    uint8_t data[105];
    uint8_t data_flag[105];
} sc_context_t;

static int channel_current;
static sc_context_t sc_context;
static TimerHandle_t timer_channel;
static TimerHandle_t timer_deal;
static xQueueHandle udp_cmd_queue = NULL;
static uint8_t broadcast_addr[6] = {0xff, 0xff, 0xff, 0xff, 0xff, 0xff};

static void channel_switch_cb(TimerHandle_t xTimer)
{
    channel_current = sc_get_next_channel(channel_current);
    wifi_mgmr_channel_set(channel_current, 1);
    // printf("[SC] switch next chan:%d\r\n", channel_current);
}

static void deal_timeout_cb(TimerHandle_t xTimer)
{
    printf("[SC] deal timeout\r\n");
}

static void wifi_event(int code, int value)
{
    switch (code)
    {
        case CODE_WIFI_ON_DISCONNECT: {
            printf("[SC] disconn, value:%d\r\n", value);
        }
        break;
        case CODE_WIFI_ON_CONNECTED: {

        }
        break;
        case CODE_WIFI_ON_GOT_IP: {
            uint16_t cmd = UDP_CMD_SEND;
            // wifi_mgmr_sta_connect_ind_stat_info_t stat;

            if (udp_cmd_queue) {
                cmd = UDP_CMD_CONNECTED;
                xQueueSend(udp_cmd_queue, &cmd, 10 / portTICK_PERIOD_MS);
                cmd = UDP_CMD_SEND;
                xQueueSend(udp_cmd_queue, &cmd, 10 / portTICK_PERIOD_MS);
            }

            // wifi_mgmr_sta_connect_ind_stat_get(&stat);
            // axk_at_wifi_sc_connect_ap_success_export(stat.ssid, stat.passphr);
        }
        break;
        default: {
            printf("[SC] Unknown code %u\r\n", code);
        }
    }
}

static uint8_t crc8_maxim(uint8_t* data, uint8_t len)
{
	uint8_t crc;
    int i;

	crc = 0x00;
	while (len--)
	{
		crc ^= *data++;
		for (i = 0; i < 8; i++) {
			if (crc & 0x01) {
				crc = (crc >> 1) ^ 0x8c;
			} else {
                crc >>= 1;
            }
		}
	}

	return crc;
}

/**
 * data format
 * EXTRA_HEAD_LEN : 5
 * [0]         total len : EXTRA_HEAD_LEN + ipLen + apPwdLen + apSsidLen
 * [1]         password len
 * [2]         ssid crc
 * [3]         bssid crc
 * [4]         total xor
 * [5~8]       ip addr(ipv4)
 * [9~pwlen]   password
 * [~ssidlen]  ssid
 * [~bssidlen] bssid
 */
static void data_parse(uint8_t data, uint8_t index)
{
    uint8_t total_len;
    uint8_t pwd_len;
    uint8_t check_len;
    char pwd[64];

    sc_context.data[index] = data;
    sc_context.data_flag[index] = 1;

    total_len = sc_context.data[0];
    pwd_len = sc_context.data[1];
    check_len = (sc_context.data_flag[1] == 0) ? total_len : (pwd_len + 9); 

    if (check_len == 0) {
        return ;
    }
    for (int i = 0; i < check_len; i++) {
        if (sc_context.data_flag[i] == 0) {
            return ;
        }
    }

    printf("[%d] recv data\r\n", xTaskGetTickCount());
    for (int i = 0; i < 9 + pwd_len; i++) {
        printf("0x%02x ", sc_context.data[i]);
    }
    printf("\r\nap info:\r\n");
    printf("ip:%d.%d.%d.%d\r\n", sc_context.data[5], sc_context.data[6], sc_context.data[7], sc_context.data[8]);
    printf("pwd:%.*s\r\n", sc_context.data[1], &sc_context.data[9]);
    printf("ssid:%s\r\n", sc_context.ssid);
    printf("bssid:" SC_ADDR_FORMATSTRING "\r\n", SC_MAC_FORMAT(sc_context.bssid));

    memset(pwd, 0, sizeof pwd);
    memcpy(pwd, &sc_context.data[9], sc_context.data[1]);
    axk_at_wifi_sc_get_info_export(sc_context.ssid, pwd);

    wifi_mgmr_sniffer_unregister(NULL);
    wifi_mgmr_sniffer_disable();
    wifi_sta_connect_sc(sc_context.ssid, pwd, channel_current);
}

static void sc_lock(int type, int guide_max, uint8_t *source, uint8_t *bssid)
{
    uint16_t offset;
    uint16_t *p_offset;
    wifi_mgmr_ap_item_t *ap_info;
    
    if (type == SC_LOCK_TODS) {
        p_offset = &sc_context.offset_tods;
    } else {
        p_offset = &sc_context.offset_frds;
    }
    offset = (uint16_t)(guide_max - 515);

    if (sc_context.offset_tods == 0 && sc_context.offset_frds == 0) {
        ap_info = sc_scan_get_ap_by_bssid(bssid);
        if (ap_info == NULL) {
            printf("[SC] not found bssid:" SC_ADDR_FORMATSTRING " \r\n", SC_MAC_FORMAT(bssid));
            return ;
        }
        /* stop channel switch */
        xTimerStop(timer_channel, 0);

        printf("[SC] source:" SC_ADDR_FORMATSTRING " \r\n", SC_MAC_FORMAT(source));
        printf("[SC] bssid:" SC_ADDR_FORMATSTRING " \r\n", SC_MAC_FORMAT(bssid));

        *p_offset = offset;
        memset(sc_context.ssid, 0, sizeof(sc_context.ssid));
        memcpy(sc_context.ssid, ap_info->ssid, ap_info->ssid_len);
        memcpy(sc_context.mac, source, 6);
        memcpy(sc_context.bssid, bssid, 6);
        sc_context.state = 1;

        /* channel correction */
        if (channel_current != ap_info->channel) {
            printf("[sc] switch to the correct channel(%d->%d)\r\n", channel_current, ap_info->channel);
            wifi_mgmr_channel_set(ap_info->channel, 1);
            channel_current = ap_info->channel;
        }
        
    } else {
        if (memcmp(sc_context.mac, source, 6) != 0) {
            printf("[SC] bssid inconsistent \r\n");
            return ;
        }
        *p_offset = offset;
    }
    printf("[SC][%d] lock type: type:%d offset:%d \r\n", xTaskGetTickCount(), type, offset);
}

static void sc_find(struct libwifi_frame *frame, int len)
{
    int cnt;
    int guide_max;
    uint16_t diff;

    if (channel_current != sc_context.channel) {
        printf("[SC]new channel:%d\r\n", channel_current);
        sc_context.channel = channel_current;
        // memset(sc_context.buffer_tods, 0, sizeof(sc_context.buffer_tods));
        // memset(sc_context.buffer_frds, 0, sizeof(sc_context.buffer_frds));
    }

    if (len > 570 && len < 640) {
        /* to_ds = 1, from_ds = 0 -> addr1-bssid add2-source add3-destination */
        if (frame->frame_control.flags.to_ds == 1 && frame->frame_control.flags.from_ds == 0 && sc_context.offset_tods == 0) {
            if (memcmp(frame->addr3, broadcast_addr, 6) == 0) {
                /* filter duplicate data */
                for (int i = 0; i < LOCK_BUFFER_SIZE; i++) {
                    if (sc_context.buffer_tods[i].len == len && 
                        *(uint16_t *)&(frame->addr2[4]) == sc_context.buffer_tods[i].mac_tail) {
                        return ;
                    }
                }

                /* match */
                cnt = 0;
                guide_max = len;
                for (int i = 0; i < LOCK_BUFFER_SIZE; i++) {
                    diff = (uint16_t)len - sc_context.buffer_tods[i].len;
                    diff += 3;
                    if (diff > 6) {
                        continue ;
                    }

                    if (*(uint16_t *)&(frame->addr2[4]) != sc_context.buffer_tods[i].mac_tail) {
                        continue ;
                    }

                    cnt++;
                    guide_max = sc_context.buffer_tods[i].len > guide_max ? sc_context.buffer_tods[i].len : guide_max;
                    if (cnt != 3) {
                        continue ;
                    }

                    /* lock channel */
                    sc_lock(SC_LOCK_TODS, guide_max, frame->addr2, frame->addr1);
                }

                /* pop data */
                for (int i = 0; i < LOCK_BUFFER_SIZE - 1; i++) {
                    sc_context.buffer_tods[i].data = sc_context.buffer_tods[i + 1].data;
                }

                /* insert data */
                sc_context.buffer_tods[LOCK_BUFFER_SIZE - 1].len = (uint16_t)len;
                sc_context.buffer_tods[LOCK_BUFFER_SIZE - 1].mac_tail = *(uint16_t *)&(frame->addr2[4]);

                printf("[R][T] push %d %x \r\n", len, sc_context.buffer_tods[LOCK_BUFFER_SIZE - 1].mac_tail);
            }
        } else if (frame->frame_control.flags.to_ds == 0 && frame->frame_control.flags.from_ds == 1 && sc_context.offset_frds == 0) {
            /* to_ds = 0, from_ds = 1 -> addr1-destination add2-bssid add3-source */
            if (memcmp(frame->addr1, broadcast_addr, 6) == 0) {
                /* filter duplicate data */
                for (int i = 0; i < LOCK_BUFFER_SIZE; i++) {
                    if (sc_context.buffer_frds[i].len == len && 
                        *(uint16_t *)&(frame->addr3[4]) == sc_context.buffer_frds[i].mac_tail) {
                        return ;
                    }
                }

                /* match */
                cnt = 0;
                guide_max = len;
                for (int i = 0; i < LOCK_BUFFER_SIZE; i++) {
                    uint16_t diff = (uint16_t)len - sc_context.buffer_frds[i].len;
                    diff += 3;
                    if (diff > 6) {
                        continue ;
                    }

                    if (*(uint16_t *)&(frame->addr3[4]) != sc_context.buffer_frds[i].mac_tail) {
                        continue ;
                    }

                    cnt++;
                    guide_max = sc_context.buffer_frds[i].len > guide_max ? sc_context.buffer_frds[i].len : guide_max;
                    if (cnt != 3) {
                        continue ;
                    }

                    /* lock channel */
                    sc_lock(SC_LOCK_FROMDS, guide_max, frame->addr3, frame->addr2);
                }

                /* pop data */
                for (int i = 0; i < LOCK_BUFFER_SIZE - 1; i++) {
                    sc_context.buffer_frds[i].data = sc_context.buffer_frds[i + 1].data;
                }

                /* insert data */
                sc_context.buffer_frds[LOCK_BUFFER_SIZE - 1].len = (uint16_t)len;
                sc_context.buffer_frds[LOCK_BUFFER_SIZE - 1].mac_tail = *(uint16_t *)&(frame->addr3[4]);

                printf("[R][F] push %d %x \r\n", len, sc_context.buffer_frds[LOCK_BUFFER_SIZE - 1].mac_tail);
            }
        }
    }
}

/**
 * one data format:(data code should have 2 to 65 data)
 * control byte       high 4 bits    low 4 bits
 * 1st 9bits:       0x0             crc(high)      data(high)
 * 2nd 9bits:       0x1                sequence header
 * 3rd 9bits:       0x0             crc(low)       data(low)
 * sequence header: 0,1,2,...
 */
static void sc_deal(struct libwifi_frame *frame, int len)
{
    uint8_t crc;
    uint8_t crc_buffer[2];
    uint16_t actual_len;
    static uint16_t buffer_tods[3];
    static uint16_t buffer_frds[3];
    
    /* to_ds = 1, from_ds = 0 -> addr1-bssid add2-source add3-destination */
    if (frame->frame_control.flags.to_ds == 1 && frame->frame_control.flags.from_ds == 0 && sc_context.offset_tods) {
        if (memcmp(frame->addr3, broadcast_addr, 6) == 0) {
            if (memcmp(frame->addr2, sc_context.mac, 6) == 0) {
                // printf("[R] len:%d type:%02x \r\n", len, pkt[0]);
                actual_len = len - sc_context.offset_tods - 40;
                if (actual_len > 360) {
                    return ;
                }
                if (sc_context.seqnum_tods == frame->seq_control.sequence_number) {
                    return ;
                }
                sc_context.seqnum_tods = frame->seq_control.sequence_number;
                if (actual_len != buffer_tods[2]) {
                    buffer_tods[0] = buffer_tods[1];
                    buffer_tods[1] = buffer_tods[2];
                    buffer_tods[2] = actual_len;
                    if (104u < (uint16_t)(buffer_tods[1] & 0xff)) {
                        return ;
                    }
                }
                if ((buffer_tods[1] & 0xff00) == 0) {
                    return ;
                }
                
                crc = (buffer_tods[0] & 0xf0) | ((buffer_tods[2] & 0xf0) >> 4);
                crc_buffer[0] = ((buffer_tods[0] & 0x0f) << 4) | (buffer_tods[2] & 0x0f);
                crc_buffer[1] = (uint8_t)(buffer_tods[1] & 0xff);
                printf("[R][T] crc data:%02x %02x %02x \r\n", buffer_tods[0], buffer_tods[1], buffer_tods[2]);
                if (crc8_maxim(crc_buffer, 2) == crc) {
                    printf("[R][T] data:%02x-%02x-%02x data:%02x index:%d \r\n", buffer_tods[0], buffer_tods[1], buffer_tods[2],
                                                                               crc_buffer[0], crc_buffer[1]);
                    data_parse(crc_buffer[0], crc_buffer[1]);
                    // sc_context.data[crc_buffer[1]] = crc_buffer[0];
                }
            }
        }
    } else if (frame->frame_control.flags.to_ds == 0 && frame->frame_control.flags.from_ds == 1 && sc_context.offset_frds) {
        /* to_ds = 0, from_ds = 1 -> addr1-destination add2-bssid add3-source */
        if (memcmp(frame->addr1, broadcast_addr, 6) == 0) {
            if (memcmp(frame->addr3, sc_context.mac, 6) == 0) {
                // printf("[R] len:%d type:%02x \r\n", len, pkt[0]);
                actual_len = len - sc_context.offset_frds - 40;
                if (actual_len > 360) {
                    return ;
                }
                if (sc_context.seqnum_frds == frame->seq_control.sequence_number) {
                    return ;
                }
                sc_context.seqnum_frds = frame->seq_control.sequence_number;
                if (actual_len != buffer_frds[2]) {
                    buffer_frds[0] = buffer_frds[1];
                    buffer_frds[1] = buffer_frds[2];
                    buffer_frds[2] = actual_len;
                    if (104u < (uint16_t)(buffer_frds[1] & 0xff)) {
                        return ;
                    }
                }
                if ((buffer_frds[1] & 0xff00) == 0) {
                    return ;
                }
                
                crc = (buffer_frds[0] & 0xf0) | ((buffer_frds[2] & 0xf0) >> 4);
                crc_buffer[0] = ((buffer_frds[0] & 0x0f) << 4) | (buffer_frds[2] & 0x0f);
                crc_buffer[1] = (uint8_t)(buffer_frds[1] & 0xff);
                printf("[R][F] fr crc data:%02x %02x %02x \r\n", buffer_frds[0], buffer_frds[1], buffer_frds[2]);
                if (crc8_maxim(crc_buffer, 2) == crc) {
                    printf("[R][F] data:%02x-%02x-%02x data:%02x index:%d \r\n", buffer_frds[0], buffer_frds[1], buffer_frds[2],
                                                                               crc_buffer[0], crc_buffer[1]);
                    data_parse(crc_buffer[0], crc_buffer[1]);
                    // sc_context.data[crc_buffer[1]] = crc_buffer[0];
                }
            }
        }
    }
}

static void sniffer_cb(void *env, uint8_t *pkt, int len)
{
    struct libwifi_frame *frame;

    frame = (struct libwifi_frame *)pkt;

    // WIFI_PKT_DATA
    if (frame->frame_control.type == 2) {
        sc_find(frame, len);
        if (sc_context.state != 0) {
            sc_deal(frame, len);
        }
    }
}

static void udp_data_packet(uint8_t *buff)
{
    uint8_t mac[6];
    uint32_t ip_u32, gw_u32, netmask_u32;

    /* mApSsid.length + mApPassword.length + 9 */
    buff[0] = sc_context.data[0];

    /* bssid */
    wifi_mgmr_sta_mac_get(mac);
    memcpy(buff + 1, mac, sizeof mac);

    /* local ip */
    wifi_mgmr_sta_ip_get(&ip_u32, &gw_u32, &netmask_u32);
    memcpy(buff + 7, &ip_u32, sizeof ip_u32);
}

static void _udp_task(void)
{
    int sockfd;
    int32_t udp_send_max = 20;
    uint16_t cmd_id;
    uint32_t target_ip;
    ip4_addr_t ip_addr;
    struct sockaddr_in remote;

    /* head(1) + mac[6] + ipv4[4] */
    uint8_t data[11];

    while (1) {
        if (xQueueReceive(udp_cmd_queue, &cmd_id, portMAX_DELAY)) {
            if (cmd_id == UDP_CMD_EXIT) {
                break;
            } else if(cmd_id == UDP_CMD_CONNECTED) {
                wifi_mgmr_sta_connect_ind_stat_info_t stat;

                wifi_mgmr_sta_connect_ind_stat_get(&stat);
                axk_at_wifi_sc_connect_ap_success_export(stat.ssid, stat.passphr);
            } else if(cmd_id == UDP_CMD_SEND) {
                printf("[SC] resp msg \r\n");

                // set addr
                target_ip = *(uint32_t *)(sc_context.data + 5);

                // packet data
                udp_data_packet(data);

                // socket create
                sockfd = socket(AF_INET, SOCK_DGRAM, 0);
                if (sockfd == -1) {
                    printf("[SC] udp socket create fail\r\n");
                    continue;
                }

                // send
                remote.sin_port = htons(18266);
                remote.sin_family = AF_INET;
                remote.sin_addr.s_addr = target_ip;

                for (int i = 0; i < udp_send_max; i++) {
                    lwip_sendto(sockfd, data, sizeof data, 0, (struct sockaddr *)&remote, sizeof remote);
                    vTaskDelay(100 / portTICK_PERIOD_MS);
                }

                close(sockfd);
            }
        }

        vTaskDelay(10 / portTICK_PERIOD_MS);
    }

    printf("[SC] exit udp task\r\n");
    xQueueReset(udp_cmd_queue);
    vTaskDelete(NULL);
}

int wifi_smartconfig_v1_start(void)
{
	udp_cmd_queue = xQueueCreate(5, sizeof(uint32_t));
    int ret = xTaskCreate(_udp_task, "udp", 512, NULL, 10, NULL);
    if (ret != pdPASS) {
        printf("[OTA] task create fail: %d\r\n", ret);
        return -1;
    }
    sc_scan_init();
    if (sc_scan_start() != 0) {
        sc_scan_deinit();
        printf("[SC] scan fail\r\n");
        return -1;
    }
    printf("[SC] scan finish\r\n");

    memset(&sc_context, 0, sizeof(sc_context));
    timer_channel = xTimerCreate("chan", 150 / portTICK_PERIOD_MS, pdTRUE, NULL, channel_switch_cb);
    timer_deal = xTimerCreate("deal", 150 / portTICK_PERIOD_MS, pdFALSE, NULL, deal_timeout_cb);

    channel_current = sc_get_next_channel(0);
    
    wifi_mgmr_sniffer_register(NULL, sniffer_cb);
    wifi_mgmr_sniffer_enable();
    wifi_mgmr_channel_set(channel_current, 1);

    /* start channel switch */
    xTimerStart(timer_channel, 0);

    wifi_event_register(wifi_event);
    
    return 0;
}

int wifi_smartconfig_v1_stop(void)
{
    uint16_t cmd = UDP_CMD_EXIT;
    if (udp_cmd_queue) {
        xQueueSend(udp_cmd_queue, &cmd, 10 / portTICK_PERIOD_MS);
        while (pdTRUE == xQueuePeek(udp_cmd_queue, &cmd, 200 / portTICK_PERIOD_MS)) {
            printf("[SC] wait for udp exit\r\n");
            vTaskDelay(100 / portTICK_PERIOD_MS);
        }
        vQueueDelete(udp_cmd_queue);
        udp_cmd_queue = NULL;
    }

    wifi_event_register(NULL);

    wifi_mgmr_sniffer_disable();
    wifi_mgmr_sniffer_unregister(NULL);

    xTimerDelete(timer_channel, 0);
    xTimerDelete(timer_deal, 0);

    sc_scan_deinit();

    return 0;
}
